package com.itx.gatewayx.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson2.JSON;
import com.itx.gatewayx.domain.AIMessage;
import com.itx.gatewayx.domain.dto.AIChatRequest;
import com.itx.gatewayx.domain.dto.AIChatStreamResponse;
import com.itx.gatewayx.service.AIService;
import com.itx.gatewayx.utils.DashScopeUtils;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Mono;

import java.util.*;

/**
 * AI聊天WebSocket处理器
 * 基于WebFlux的WebSocket实现
 */
@Slf4j
@Component
public class AIWebSocketHandler implements WebSocketHandler {

    @Autowired
    private AIService aiService;
    
    @Autowired
    private DashScopeUtils dashScopeUtils;
    
    @Value("${ai.dashscope.api-key}")
    private String apiKey;
    
    @Value("${ai.dashscope.model}")
    private String model;
    
    @Override
    public Mono<Void> handle(WebSocketSession session) {
        return session.receive()
                .map(message -> message.getPayloadAsText())
                .flatMap(payload -> handleMessage(payload, session))
                .then();
    }
    
    /**
     * 处理客户端发来的消息
     */
    private Mono<Void> handleMessage(String message, WebSocketSession session) {
        return Mono.fromCallable(() -> {
            try {
                // 解析客户端消息
                AIChatRequest request = JSON.parseObject(message, AIChatRequest.class);
                
                // 获取用户ID（如果有）
                Long userId = null;
                if (request.getToken() != null && !request.getToken().isEmpty()) {
                    // 这里可以根据token获取用户ID
                    // userId = UserUtils.getUserIdFromToken(request.getToken());
                }
                
                // 准备消息上下文
                List<AIMessage> contextMessages = new ArrayList<>();
                
                // 优先使用客户端传递的历史消息，如果为空则从数据库获取
                if (request.getHistory() != null && !request.getHistory().isEmpty()) {
                    contextMessages.addAll(request.getHistory());
                }
                // 如果客户端没有传递历史消息，且有会话ID，则从数据库获取
                else if (request.getConversationId() != null) {
                    contextMessages = aiService.getMessagesByConversationId(request.getConversationId(), userId);
                }
                
                // 添加当前用户消息
                AIMessage userMessage = AIMessage.builder()
                        .role("user")
                        .content(request.getMessage())
                        .build();
                contextMessages.add(userMessage);
                
                // 发送开始消息
                AIChatStreamResponse startResponse = AIChatStreamResponse.builder()
                        .type("start")
                        .messageId(UUID.randomUUID().toString())
                        .conversationId(request.getConversationId())
                        .build();
                
                session.send(Mono.just(session.textMessage(JSON.toJSONString(startResponse)))).subscribe();
                
                // 调用流式API处理
                streamChatToClient(session, contextMessages, request);
                
                return true;
            } catch (Exception e) {
                log.error("处理WebSocket消息失败", e);
                // 发送错误消息
                AIChatStreamResponse errorResponse = AIChatStreamResponse.builder()
                        .type("error")
                        .messageId(UUID.randomUUID().toString())
                        .content("处理消息失败: " + e.getMessage())
                        .build();
                
                session.send(Mono.just(session.textMessage(JSON.toJSONString(errorResponse)))).subscribe();
                return false;
            }
        }).flatMap(result -> Mono.empty());
    }
    
    /**
     * 流式调用AI聊天并实时发送结果给客户端
     */
    private void streamChatToClient(WebSocketSession session, List<AIMessage> contextMessages, AIChatRequest request) {
        try {
            // 生成消息ID
            String messageId = UUID.randomUUID().toString();
            
            // 转换消息格式
            List<Message> dashScopeMessages = new ArrayList<>();
            for (AIMessage message : contextMessages) {
                Role role;
                if ("user".equals(message.getRole())) {
                    role = Role.USER;
                } else if ("assistant".equals(message.getRole())) {
                    role = Role.ASSISTANT;
                } else {
                    role = Role.SYSTEM;
                }
                
                Message dashScopeMessage = Message.builder()
                        .role(role.getValue())
                        .content(message.getContent())
                        .build();
                
                dashScopeMessages.add(dashScopeMessage);
            }
            
            // 构建参数
            GenerationParam param = GenerationParam.builder()
                    .apiKey(apiKey)
                    .model(model)
                    .messages(dashScopeMessages)
                    .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                    .incrementalOutput(true)
                    .build();
            
            // 创建StringBuilder收集完整回复
            StringBuilder finalContent = new StringBuilder();
            
            // 调用流式API
            Generation generation = new Generation();
            Flowable<GenerationResult> result = generation.streamCall(param);
            
            // 处理流式响应
            result.subscribe(
                // onNext - 处理每个流式响应片段
                message -> {
                    try {
                        String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                        
                        if (content != null && !content.isEmpty()) {
                            // 累积内容
                            finalContent.append(content);
                            
                            // 发送流式响应
                            AIChatStreamResponse streamResponse = AIChatStreamResponse.builder()
                                    .type("chunk")
                                    .messageId(messageId)
                                    .conversationId(request.getConversationId())
                                    .content(content)
                                    .build();
                            
                            session.send(Mono.just(session.textMessage(JSON.toJSONString(streamResponse)))).subscribe();
                        }
                    } catch (Exception e) {
                        log.error("处理流式响应片段失败", e);
                    }
                },
                // onError - 处理错误
                error -> {
                    log.error("流式调用AI服务失败", error);
                    
                    // 发送错误消息
                    AIChatStreamResponse errorResponse = AIChatStreamResponse.builder()
                            .type("error")
                            .messageId(messageId)
                            .conversationId(request.getConversationId())
                            .content("AI服务调用失败: " + error.getMessage())
                            .build();
                    
                    session.send(Mono.just(session.textMessage(JSON.toJSONString(errorResponse)))).subscribe();
                },
                // onComplete - 处理完成
                () -> {
                    log.info("流式调用AI服务完成");
                    
                    try {
                        // 保存消息到数据库
                        Long userId = null;
                        if (request.getToken() != null && !request.getToken().isEmpty()) {
                            // userId = UserUtils.getUserIdFromToken(request.getToken());
                        }
                        
                        // 保存用户消息和AI回复 (非阻塞方式处理)
                        Mono.fromRunnable(() -> {
                            try {
                                aiService.sendMessage(userId, request);
                            } catch (Exception e) {
                                log.error("保存消息失败", e);
                            }
                        }).subscribe();
                        
                        // 发送完成消息
                        AIChatStreamResponse completeResponse = AIChatStreamResponse.builder()
                                .type("end")
                                .messageId(messageId)
                                .conversationId(request.getConversationId())
                                .content(finalContent.toString())
                                .build();
                        
                        session.send(Mono.just(session.textMessage(JSON.toJSONString(completeResponse)))).subscribe();
                    } catch (Exception e) {
                        log.error("完成流式处理失败", e);
                    }
                }
            );
            
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            log.error("初始化流式调用AI服务失败", e);
            
            // 发送错误消息
            AIChatStreamResponse errorResponse = AIChatStreamResponse.builder()
                    .type("error")
                    .messageId(UUID.randomUUID().toString())
                    .conversationId(request.getConversationId())
                    .content("AI服务调用失败: " + e.getMessage())
                    .build();
            
            session.send(Mono.just(session.textMessage(JSON.toJSONString(errorResponse)))).subscribe();
        }
    }
} 